from langchain.prompts import PromptTemplate
from models.custom_llm import CustomLLM
# from ChatGLM2_6B.test_langchain import create_model
from test_langchain import create_model
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)


# llm = CustomLLM()
llm = create_model(modelpath="THUDM/chatglm2-6b")

prompt = PromptTemplate(
    input_variables=["product"],
    template="请给一个制造{product}的公司起一个好听的名字",
)

from langchain.chains import LLMChain
chain = LLMChain(llm=llm, prompt=prompt)
# print(chain.run("彩虹色的袜子"))
print(llm.predict("what is the date today"))
print(llm.predict_messages([HumanMessage(content="Translate this sentence from English to French. I love programming.")]))